#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import uuid
import time
import tempfile
import shutil
import subprocess
import requests
import m3u8
from enum import IntEnum
from urllib.parse import urljoin
from concurrent.futures import ThreadPoolExecutor, as_completed
from PySide6.QtCore import QRunnable, QObject, Signal, Slot


class DownloadStatus(IntEnum):
    """下载状态枚举"""
    WAITING = 0      # 等待中
    DOWNLOADING = 1  # 下载中
    PAUSED = 2       # 已暂停  
    COMPLETED = 3    # 已完成
    FAILED = 4       # 失败


class DownloadTask(QRunnable):
    """M3U8下载任务"""
    
    def __init__(self, task_id, url, save_path, signals=None):
        super().__init__()
        self.task_id = task_id
        self.url = url
        self.save_path = save_path
        self.signals = signals
        self.status = DownloadStatus.WAITING
        self.progress = 0
        self.temp_dir = None
        self.ts_files = []
        self.downloaded_files = []
        self.max_workers = 10  # 单个任务的并发数
        self.should_stop = False
    
    def run(self):
        """执行下载任务"""
        try:
            self._update_status(DownloadStatus.DOWNLOADING)
            
            # 创建临时目录
            self.temp_dir = tempfile.mkdtemp()
            
            # 下载并解析m3u8文件
            playlist = self._download_m3u8()
            
            # 如果是master playlist，选择第一个stream
            if playlist.is_variant:
                if not playlist.playlists:
                    raise Exception("未找到可用的播放列表")
                
                # 选择画质最高的流
                playlist.playlists.sort(key=lambda p: p.stream_info.bandwidth if p.stream_info else 0, reverse=True)
                stream_url = urljoin(self.url, playlist.playlists[0].uri)
                playlist = self._download_m3u8(stream_url)
            
            # 获取所有ts文件
            base_uri = playlist.base_uri or os.path.dirname(self.url)
            self.ts_files = [urljoin(base_uri, segment.uri) for segment in playlist.segments]
            total_files = len(self.ts_files)
            
            if total_files == 0:
                raise Exception("未找到任何媒体片段")
            
            # 优化：使用as_completed来跟踪完成情况
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                # 准备下载任务
                future_to_index = {}
                for i, ts_url in enumerate(self.ts_files):
                    if self.should_stop:
                        break
                    
                    # 设置保存路径
                    filename = f"{i:05d}.ts"
                    save_path = os.path.join(self.temp_dir, filename)
                    
                    # 提交下载任务
                    future = executor.submit(self._download_file, ts_url, save_path, i)
                    future_to_index[future] = i
                
                # 使用as_completed跟踪任务完成情况
                completed_count = 0
                for future in as_completed(future_to_index):
                    if self.should_stop:
                        break
                    
                    result, index, filepath = future.result()
                    if result:
                        completed_count += 1
                        self.progress = int(completed_count / total_files * 100)
                        self._update_progress(self.progress)
                        self.downloaded_files.append(filepath)
            
            # 检查是否被停止
            if self.should_stop:
                self._update_status(DownloadStatus.PAUSED)
                return
            
            # 确保所有文件都已下载
            if len(self.downloaded_files) != total_files:
                raise Exception(f"下载不完整：预期 {total_files} 个文件，实际下载 {len(self.downloaded_files)} 个")
            
            # 合并文件
            output_file = self._get_output_filename()
            self._merge_files(output_file)
            
            # 清理临时文件
            self._cleanup()
            
            # 更新状态
            self._update_status(DownloadStatus.COMPLETED)
            self._update_progress(100)
            
            # 触发完成信号
            if self.signals:
                self.signals.task_completed.emit(self.task_id)
        
        except Exception as e:
            self._update_status(DownloadStatus.FAILED)
            if self.signals:
                self.signals.task_failed.emit(self.task_id, str(e))
            print(f"任务 {self.task_id} 失败: {str(e)}")
            self._cleanup()
    
    def _download_m3u8(self, url=None):
        """下载并解析m3u8文件"""
        target_url = url or self.url
        try:
            response = requests.get(target_url, timeout=30)
            response.raise_for_status()
            
            return m3u8.loads(response.text, uri=target_url)
        except requests.RequestException as e:
            raise Exception(f"下载M3U8文件失败: {str(e)}")
    
    def _download_file(self, url, save_path, index):
        """
        下载单个文件
        
        Args:
            url: 文件URL
            save_path: 保存路径
            index: 文件索引
            
        Returns:
            tuple: (是否成功, 文件索引, 文件路径)
        """
        try:
            # 如果已暂停，直接返回失败
            if self.should_stop:
                return False, index, save_path
                
            response = requests.get(url, stream=True, timeout=30)
            response.raise_for_status()
            
            with open(save_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if self.should_stop:
                        return False, index, save_path
                    if chunk:
                        f.write(chunk)
            
            return True, index, save_path
        except Exception as e:
            print(f"下载文件失败 {url}: {str(e)}")
            return False, index, save_path
    
    def _get_output_filename(self):
        """获取输出文件名"""
        base_name = os.path.basename(self.url).split('?')[0]
        if not base_name or base_name.endswith('.m3u8'):
            base_name = f"video_{int(time.time())}.mp4"
        elif not base_name.endswith('.mp4'):
            base_name = f"{base_name}.mp4"
        
        return os.path.join(self.save_path, base_name)
    
    def _merge_files(self, output_file):
        """合并ts文件为mp4"""
        # 创建保存目录（如果不存在）
        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        
        # 创建文件列表
        list_file = os.path.join(self.temp_dir, "filelist.txt")
        with open(list_file, 'w') as f:
            # 确保按文件名的数字顺序排序
            ts_files = sorted([f for f in os.listdir(self.temp_dir) if f.endswith('.ts')], 
                             key=lambda x: int(x.split('.')[0]))
            for ts_file in ts_files:
                f.write(f"file '{os.path.join(self.temp_dir, ts_file)}'\n")
        
        # 使用ffmpeg合并 - 异步执行并捕获输出
        try:
            subprocess.run([
                'ffmpeg', '-f', 'concat', '-safe', '0', 
                '-i', list_file, '-c', 'copy', output_file,
                '-loglevel', 'warning'  # 降低日志级别，减少输出
            ], check=True, capture_output=True)
        except subprocess.CalledProcessError as e:
            raise Exception(f"合并文件失败: {e.stderr.decode('utf-8', errors='ignore')}")
        except FileNotFoundError:
            raise Exception("ffmpeg未安装，请先安装ffmpeg")
    
    def _cleanup(self):
        """清理临时文件"""
        if self.temp_dir and os.path.exists(self.temp_dir):
            try:
                shutil.rmtree(self.temp_dir)
                self.temp_dir = None
            except Exception as e:
                print(f"清理临时文件失败: {str(e)}")
    
    def _update_status(self, status):
        """更新状态"""
        self.status = status
        if self.signals:
            self.signals.status_updated.emit(self.task_id, status)
    
    def _update_progress(self, progress):
        """更新进度"""
        self.progress = progress
        if self.signals:
            self.signals.progress_updated.emit(self.task_id, progress)
    
    def pause(self):
        """暂停任务"""
        self.should_stop = True
        self._update_status(DownloadStatus.PAUSED)
    
    def resume(self):
        """恢复任务"""
        self.should_stop = False 